#!/bin/bash

# ==============================================
# WSL USB 设备绑定辅助脚本
# 用于在 WSL 环境中自动检测和提示 USB 设备绑定
# ==============================================

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
MAGENTA='\033[0;35m'
BOLD='\033[1m'
NC='\033[0m' # No Color

info() { echo -e "${GREEN}[INFO]${NC} $1"; }
warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1"; }
step() { echo -e "${MAGENTA}${BOLD}=== $1 ===${NC}"; }
cmd() { echo -e "${CYAN}$1${NC}"; }

# 检测是否在 WSL 环境
check_wsl() {
    if grep -qi microsoft /proc/version; then
        return 0  # 是 WSL
    else
        return 1  # 不是 WSL
    fi
}

# 显示 WSL USB 设备绑定指南
show_wsl_usb_guide() {
    echo ""
    echo "${BOLD}${BLUE}╔════════════════════════════════════════════════════════════════╗${NC}"
    echo "${BOLD}${BLUE}║         WSL2 USB 设备绑定指南                                 ║${NC}"
    echo "${BOLD}${BLUE}╚════════════════════════════════════════════════════════════════╝${NC}"
    echo ""
    
    step "检测到您在 WSL 环境中运行"
    echo ""
    
    info "WSL2 默认不支持直接访问 USB 设备，需要通过 usbipd-win 转发"
    echo ""
    
    echo "${BOLD}📋 操作步骤：${NC}"
    echo ""
    
    echo "${BOLD}${GREEN}【步骤 1】${NC} 在 Windows 主机上安装 usbipd-win"
    echo "   打开 PowerShell (管理员权限)，运行："
    cmd "   winget install --interactive --exact dorssel.usbipd-win"
    echo "   或从官网下载: ${CYAN}https://github.com/dorssel/usbipd-win/releases${NC}"
    echo ""
    
    echo "${BOLD}${GREEN}【步骤 2】${NC} 在 WSL 中安装 USB/IP 客户端工具"
    echo "   如果还没安装，运行："
    cmd "   sudo apt update"
    cmd "   sudo apt install linux-tools-generic hwdata"
    cmd "   sudo update-alternatives --install /usr/local/bin/usbip usbip /usr/lib/linux-tools/*-generic/usbip 20"
    echo ""
    
    echo "${BOLD}${GREEN}【步骤 3】${NC} 将 Orange Pi 5 Plus 连接到电脑并进入 Maskrom 模式"
    echo "   1. 关闭 Orange Pi 5 Plus 电源"
    echo "   2. 按住 ${YELLOW}Maskrom 按钮${NC}（靠近 USB-C 接口）"
    echo "   3. 连接 USB-C 数据线到电脑"
    echo "   4. 等待 2 秒后松开按钮"
    echo ""
    
    echo "${BOLD}${GREEN}【步骤 4】${NC} 在 Windows PowerShell (管理员) 中查找设备"
    cmd "   usbipd list"
    echo ""
    echo "   查找类似以下输出的设备："
    echo "   ${CYAN}BUSID  VID:PID    DEVICE${NC}"
    echo "   ${CYAN}1-4    2207:350b  Rockchip Maskrom Device${NC}"
    echo ""
    echo "   ${YELLOW}记下 BUSID (例如: 1-4)${NC}"
    echo ""
    
    echo "${BOLD}${GREEN}【步骤 5】${NC} 绑定设备到 WSL"
    echo "   在 PowerShell (管理员) 中运行："
    cmd "   usbipd bind --busid 1-4"
    echo "   ${YELLOW}(将 1-4 替换为您实际的 BUSID)${NC}"
    echo ""
    
    echo "${BOLD}${GREEN}【步骤 6】${NC} 附加设备到 WSL"
    echo "   在 PowerShell (管理员) 中运行："
    cmd "   usbipd attach --wsl --busid 1-4"
    echo "   ${YELLOW}(将 1-4 替换为您实际的 BUSID)${NC}"
    echo ""
    
    echo "${BOLD}${GREEN}【步骤 7】${NC} 在 WSL 中验证设备"
    echo "   回到 WSL 终端，运行："
    cmd "   lsusb"
    echo ""
    echo "   应该看到类似输出："
    echo "   ${CYAN}Bus 001 Device 002: ID 2207:350b Fuzhou Rockchip Electronics Company${NC}"
    echo ""
    cmd "   sudo rkdeveloptool ld"
    echo ""
    echo "   应该看到："
    echo "   ${CYAN}DevNo=1 Vid=0x2207,Pid=0x350b ... Mode=Maskrom${NC}"
    echo ""
    
    echo "${BOLD}💡 提示：${NC}"
    echo "   • 每次重新连接设备都需要重新执行 ${YELLOW}步骤 6${NC} (attach)"
    echo "   • 如果设备断开，运行: ${CYAN}usbipd detach --busid 1-4${NC}"
    echo "   • 可以保持 PowerShell 窗口打开以便快速重新绑定"
    echo ""
    
    echo "${BOLD}🚀 完成绑定后，继续运行：${NC}"
    cmd "   sudo bash make_flash.sh target=SD uimg=xxx rootfs=xxx"
    echo ""
    
    echo "${BOLD}${BLUE}════════════════════════════════════════════════════════════════${NC}"
}

# 检查设备是否已连接
check_device() {
    info "检查 Rockchip 设备..."
    
    if lsusb | grep -i "2207:350b" > /dev/null 2>&1; then
        echo ""
        echo "${GREEN}✓${NC} 检测到 Rockchip Maskrom 设备！"
        lsusb | grep -i "2207"
        echo ""
        
        if command -v rkdeveloptool &> /dev/null; then
            if sudo rkdeveloptool ld 2>/dev/null | grep -q "Maskrom"; then
                echo "${GREEN}✓${NC} rkdeveloptool 可以访问设备"
                sudo rkdeveloptool ld
                echo ""
                echo "${GREEN}${BOLD}✓✓✓ 设备就绪，可以开始刷写！${NC}"
                return 0
            else
                warn "设备已连接但 rkdeveloptool 无法访问"
                warn "可能需要 sudo 权限或重新插拔设备"
            fi
        else
            warn "rkdeveloptool 未安装"
            info "安装方法见文档"
        fi
    else
        echo ""
        warn "未检测到 Rockchip 设备"
        echo ""
        echo "请确保："
        echo "  1. 设备已进入 Maskrom 模式"
        echo "  2. USB 已通过 usbipd attach 绑定到 WSL"
        echo ""
    fi
    
    return 1
}

# 显示快速参考卡片
show_quick_reference() {
    echo ""
    echo "${BOLD}${CYAN}┌─────────────────────────────────────────────────┐${NC}"
    echo "${BOLD}${CYAN}│  WSL USB 快速参考                              │${NC}"
    echo "${BOLD}${CYAN}├─────────────────────────────────────────────────┤${NC}"
    echo "${CYAN}│${NC} ${BOLD}Windows PowerShell (管理员):${NC}"
    echo "${CYAN}│${NC}   查看设备: ${YELLOW}usbipd list${NC}"
    echo "${CYAN}│${NC}   绑定设备: ${YELLOW}usbipd bind --busid 1-4${NC}"
    echo "${CYAN}│${NC}   附加设备: ${YELLOW}usbipd attach --wsl --busid 1-4${NC}"
    echo "${CYAN}│${NC}   分离设备: ${YELLOW}usbipd detach --busid 1-4${NC}"
    echo "${CYAN}│${NC}"
    echo "${CYAN}│${NC} ${BOLD}WSL 终端:${NC}"
    echo "${CYAN}│${NC}   查看设备: ${YELLOW}lsusb${NC}"
    echo "${CYAN}│${NC}   验证连接: ${YELLOW}sudo rkdeveloptool ld${NC}"
    echo "${CYAN}│${NC}   开始刷写: ${YELLOW}sudo bash make_flash.sh ...${NC}"
    echo "${BOLD}${CYAN}└─────────────────────────────────────────────────┘${NC}"
    echo ""
}

# 创建 Windows 端的辅助脚本
create_windows_script() {
    local SCRIPT_NAME="usb_attach.ps1"
    
    cat > "$SCRIPT_NAME" << 'EOF'
# WSL USB 自动绑定脚本
# 在 Windows PowerShell (管理员) 中运行

Write-Host "正在查找 Rockchip 设备..." -ForegroundColor Green

# 查找设备
$devices = usbipd list | Select-String "2207:350b"

if ($devices) {
    Write-Host "找到 Rockchip 设备:" -ForegroundColor Green
    Write-Host $devices
    
    # 提取 BUSID
    $busid = ($devices -split '\s+')[0]
    Write-Host "BUSID: $busid" -ForegroundColor Yellow
    
    # 询问是否绑定
    $confirm = Read-Host "是否绑定并附加到 WSL? (y/n)"
    
    if ($confirm -eq 'y' -or $confirm -eq 'Y') {
        Write-Host "绑定设备..." -ForegroundColor Green
        usbipd bind --busid $busid
        
        Write-Host "附加到 WSL..." -ForegroundColor Green
        usbipd attach --wsl --busid $busid
        
        Write-Host "完成! 现在可以在 WSL 中使用设备了" -ForegroundColor Green
        Write-Host "在 WSL 中运行: lsusb" -ForegroundColor Cyan
    }
} else {
    Write-Host "未找到 Rockchip 设备" -ForegroundColor Red
    Write-Host "请确保设备已连接并进入 Maskrom 模式" -ForegroundColor Yellow
    Write-Host "" 
    Write-Host "所有 USB 设备:" -ForegroundColor Cyan
    usbipd list
}

Read-Host "按回车键退出"
EOF
    
    info "已创建 Windows 辅助脚本: ${CYAN}$SCRIPT_NAME${NC}"
    echo "   将此文件复制到 Windows 系统，用管理员权限运行"
}

# 主函数
main() {
    clear
    echo ""
    echo "${BOLD}${BLUE}════════════════════════════════════════════════════════════════${NC}"
    echo "${BOLD}${BLUE}   WSL USB 设备绑定辅助工具${NC}"
    echo "${BOLD}${BLUE}   Orange Pi 5 Plus Flashing Tool for WSL${NC}"
    echo "${BOLD}${BLUE}════════════════════════════════════════════════════════════════${NC}"
    echo ""
    
    if ! check_wsl; then
        info "您不在 WSL 环境中，USB 设备应该可以直接访问"
        exit 0
    fi
    
    # 显示完整指南
    show_wsl_usb_guide
    
    # 检查设备状态
    if check_device; then
        echo "${GREEN}${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
        echo "${GREEN}${BOLD}  ✓ 一切就绪！可以开始刷写系统了${NC}"
        echo "${GREEN}${BOLD}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
        echo ""
        exit 0
    fi
    
    # 显示快速参考
    show_quick_reference
    
    # 询问是否创建 Windows 脚本
    echo ""
    read -p "是否创建 Windows 端辅助脚本? (y/n): " create_script
    if [[ "$create_script" =~ ^[Yy]$ ]]; then
        create_windows_script
        echo ""
        info "脚本已创建，您可以通过以下方式使用："
        echo "   1. 复制 ${CYAN}usb_attach.ps1${NC} 到 Windows 系统"
        echo "   2. 右键 -> ${YELLOW}使用 PowerShell 运行${NC} (管理员)"
        echo "   3. 或在 PowerShell (管理员) 中: ${CYAN}./usb_attach.ps1${NC}"
    fi
    
    echo ""
    warn "完成 USB 绑定后，重新运行此脚本验证连接"
    echo ""
}

# 如果有参数，执行快速检查
if [ "$1" = "check" ]; then
    check_wsl && check_device
    exit $?
fi

# 执行主函数
main "$@"

